///  Copyright (c) Microsoft Corporation.  All rights reserved.
///
///     Sample code co-developed by William Dixon of V6 Security, Inc.
///     (http://www.v6security.com) and Dan Griffin of JW Secure, Inc.
///     Errata will be published via http://www.jwsecure.com/dan.
///

#ifndef UNICODE
#define UNICODE
#endif

#ifndef WINVER
#define WINVER 0x0600
#endif

#ifndef _WIN32_WINNT
#define _WIN32_WINNT 0x0600
#endif

#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif

#ifdef IPV6STRICT
#undef IPV6STRICT
#endif

#include <rpc.h>
#include <ntdsapi.h>
#include <wchar.h>
#include <winsock2.h>
#include <ws2ipdef.h>
#include <mstcpip.h>
#include <ws2tcpip.h>

#include <windows.h>
#include <fwpmu.h>
#include <ipsectypes.h>

#include <stdio.h>
#include <strsafe.h>
#include <wincrypt.h>
#include <inttypes.h>

#ifndef IPV6STRICT
#define IPV6STRICT
#endif

#define cCANDIDATE_FILTERS_REQUESTED                    256

//
// Heap helpers
//

LPVOID Alloc(DWORD cb)
{
    return HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, cb);
}

void Free(LPVOID pv)
{
    HeapFree(GetProcessHeap(), 0, pv);
}

//
// Flow macros
//

static 
void
WINAPI
_OutputDbgStr(
    __in  LPSTR szMsg,
    __in  DWORD dwStatus)
{
    CHAR rgsz [256];

    StringCbPrintfA(
        rgsz,
        sizeof(rgsz),
        "ERROR: %s - 0x%x\n",
        szMsg,
        dwStatus);
    OutputDebugStringA(rgsz);
}

#define TRY_DWORD(_X) {                                             \
    if (ERROR_SUCCESS != (status = _X)) {                           \
        _OutputDbgStr(#_X, status);                                 \
        __leave;                                                    \
    }                                                               \
}

#define TRY_BOOL(_X) {                                              \
    if (FALSE == (_X)) {                                            \
        status = GetLastError();                                    \
        _OutputDbgStr(#_X, status);                                 \
        __leave;                                                    \
    }                                                               \
}

#define TRY_ALLOC(_X) {                                             \
    if (NULL == (_X)) {                                             \
        status = ERROR_NOT_ENOUGH_MEMORY;                           \
        __leave;                                                    \
    }                                                               \
}

#define TRY_WSACALL(_X) {                                           \
    if (ERROR_SUCCESS != (_X)) {                                    \
        status = WSAGetLastError();                                 \
        _OutputDbgStr(#_X, status);                                 \
        __leave;                                                    \
    }                                                               \
}

VOID
DisplayIPv4AddressPair(
    __in const unsigned int localV4Address, 
    __in const unsigned int remoteV4Address)
{
    wprintf(L", %ld.%ld.%ld.%ld, %ld.%ld.%ld.%ld",
        ((byte*)&localV4Address)[3],
        ((byte*)&localV4Address)[2],
        ((byte*)&localV4Address)[1],
        ((byte*)&localV4Address)[0],
        ((byte*)&remoteV4Address)[3],
        ((byte*)&remoteV4Address)[2],
        ((byte*)&remoteV4Address)[1],
        ((byte*)&remoteV4Address)[0]);
}

VOID
DisplayIPv6AddressPair(
    __in const unsigned char localV6Address[16],
    __in const unsigned char remoteV6Address[16])
{
    wprintf(L", [%x:%x:%x:%x%x:%x:%x:%x], [%x:%x:%x:%x%x:%x:%x:%x]",
        ((WORD*)&localV6Address)[7],
        ((WORD*)&localV6Address)[6],
        ((WORD*)&localV6Address)[5],
        ((WORD*)&localV6Address)[4],
        ((WORD*)&localV6Address)[3],
        ((WORD*)&localV6Address)[2],
        ((WORD*)&localV6Address)[1],
        ((WORD*)&localV6Address)[0],
        ((WORD*)&remoteV6Address)[7],
        ((WORD*)&remoteV6Address)[6],
        ((WORD*)&remoteV6Address)[5],
        ((WORD*)&remoteV6Address)[4],
        ((WORD*)&remoteV6Address)[3],
        ((WORD*)&remoteV6Address)[2],
        ((WORD*)&remoteV6Address)[1],
        ((WORD*)&remoteV6Address)[0]);
}

VOID
DisplayByteBlob(
    __in    FWP_BYTE_BLOB *pBlob)
{
    DWORD iByte = 0;
    CHAR rgbRow [256];

    if (NULL == pBlob)
        return;

    while (iByte < pBlob->size)
    {
        StringCbPrintfA(
            rgbRow + (iByte * 2), 
            sizeof(rgbRow) - (iByte * 2),
            "%02X", 
            pBlob->data [iByte]);
        iByte++;
    }

    wprintf(L" [ %S ]", rgbRow);
}

VOID
DisplayCertInfo(
    __in    IKEEXT_CERTIFICATE_CREDENTIAL0 *pCert)
{
    DisplayByteBlob(&pCert->certHash);
}

VOID
DisplayPresharedKeyInfo(
    __in    IKEEXT_PRESHARED_KEY_AUTHENTICATION0 *pKey)
{
    DisplayByteBlob(&pKey->presharedKey);
}

//
// Display quick-mode security association information to the console
// 
VOID
PrintQmSa(
    __in  const IPSEC_SA_CONTEXT0* qmSa)
{
    UINT32 i = 0;
    IPSEC_AUTH_TRANSFORM_ID0* authTransformId = NULL;
    IPSEC_CIPHER_TRANSFORM_ID0* cipherTransformId = NULL;
    IPSEC_SA0* sa = NULL;
    BOOL bLogged = FALSE;

    // Static array containing string description of various encryption 
    // algorithms
    static const wchar_t* const cipherString[] =
    {
        L"None",
        L"DES",
        L"3DES",
        L"AES-128",
        L"AES-192",
        L"AES-256"
    };
    // Static array containing string description of various hashing algorithms
    static const wchar_t* const hashString[] =
    {
        L"MD5",
        L"SHA1"
    };

    // Static array containing string description of various traffic types
    static const wchar_t* const trafficString[] =
    {
        L"transport",
        L"tunnel"
    };

    wprintf(L"IKEQM: ");
    if(qmSa->outboundSa->traffic.trafficType < RTL_NUMBER_OF(trafficString))
    {
        wprintf(L"%s", trafficString[(int)qmSa->outboundSa->traffic.trafficType]);
        bLogged = TRUE;
    }
    //-----------------------------------------
    // Print all the SAs in the bundle
    for(i=0; i<qmSa->outboundSa->saBundle.numSAs; i++)
    {
        sa = &qmSa->outboundSa->saBundle.saList[i];

        // For each SA print information based on its type
        switch(sa->saTransformType)
        {
            case IPSEC_TRANSFORM_AH:
            {
                // Print AH hash algorithm
                authTransformId = 
                    &sa->ahInformation->authTransform.authTransformId;
                if(authTransformId->authType < RTL_NUMBER_OF(hashString))
                {
                    if(TRUE == bLogged)
                    {
                        wprintf(L", %s", hashString[authTransformId->authType]);
                    }
                    else
                    {
                        wprintf(L"%s", hashString[authTransformId->authType]);
                        bLogged = TRUE;
                    }
                }
                break;
            }
            case IPSEC_TRANSFORM_ESP_CIPHER:
            {
                // Print ESP encryption algorithm
                cipherTransformId = 
                    &sa->espCipherInformation->cipherTransform.cipherTransformId;
                if(cipherTransformId->cipherType < RTL_NUMBER_OF(cipherString))
                {
                    if(TRUE == bLogged)
                    {
                        wprintf(L", %s", cipherString[cipherTransformId->cipherType]);
                    }
                    else
                    {
                        wprintf(L"%s", cipherString[cipherTransformId->cipherType]);
                        bLogged = TRUE;
                    }
                }
                break;
            }
            case IPSEC_TRANSFORM_ESP_AUTH:
            {
                // Print ESP hash algorithm
                authTransformId = 
                    &sa->espAuthInformation->authTransform.authTransformId;
                if(authTransformId->authType < RTL_NUMBER_OF(hashString))
                {
                    if(TRUE == bLogged)
                    {
                        wprintf(L", %s", hashString[authTransformId->authType]);
                    }
                    else
                    {
                        wprintf(L"%s", hashString[authTransformId->authType]);
                        bLogged = TRUE;
                    }
                }
                break;
            }
            case IPSEC_TRANSFORM_ESP_AUTH_AND_CIPHER:
            {
                // Print ESP hash & encryption algorithm
                authTransformId = 
                    &sa->espAuthAndCipherInformation->saAuthInformation.authTransform.authTransformId;
                cipherTransformId = 
                    &sa->espAuthAndCipherInformation->saCipherInformation.cipherTransform.cipherTransformId;
                if((authTransformId->authType < RTL_NUMBER_OF(hashString)) &&
                (cipherTransformId->cipherType < RTL_NUMBER_OF(cipherString)))
                {
                    if(TRUE == bLogged)
                    {
                        wprintf(L", %s, %s",
                            hashString[authTransformId->authType],
                            cipherString[cipherTransformId->cipherType]);
                    }
                    else
                    {
                        wprintf(L"%s, %s",
                            hashString[authTransformId->authType],
                            cipherString[cipherTransformId->cipherType]);
                        bLogged = TRUE;
                    }
                }
                break;
            }
            default:
            {
                break;
            }
        }
    }
    if(TRUE == bLogged)
    {
        wprintf(L", %ldkb, %lds",  
            qmSa->outboundSa->saBundle.lifetime.lifetimeKilobytes, 
            qmSa->outboundSa->saBundle.lifetime.lifetimeSeconds);
    }
    else
    {
        wprintf(L"%ldkb, %lds",  
            qmSa->outboundSa->saBundle.lifetime.lifetimeKilobytes, 
            qmSa->outboundSa->saBundle.lifetime.lifetimeSeconds);
    }
    
    //at this moment bLogged is true.
    switch(qmSa->outboundSa->traffic.ipVersion)
    {
    case FWP_IP_VERSION_V4:
        DisplayIPv4AddressPair(
            qmSa->outboundSa->traffic.localV4Address,
            qmSa->outboundSa->traffic.remoteV4Address);
         break;
    case FWP_IP_VERSION_V6:
        DisplayIPv6AddressPair(
            qmSa->outboundSa->traffic.localV6Address,
            qmSa->outboundSa->traffic.remoteV6Address);
        break;
    default:
        break;
    }

    wprintf(L", idletime=%lds", qmSa->outboundSa->saBundle.idleTimeoutSeconds);
    wprintf(L"\n");
}

//
// Display quick-mode security association information to the console
// 
VOID
PrintMmSa(
   __in const IKEEXT_SA_DETAILS0* mmSa)
{
    BOOL bLogged = FALSE;
    // Static array containing string description of various encryption 
    // algorithms
    static const wchar_t* const cipherString[] =
    {
        L"DES",
        L"3DES",
        L"AES-128",
        L"AES-192",
        L"AES-256"
    };
    // Static array containing string description of various hashing algorithms
    static const wchar_t* const hashString[] =
    {
        L"MD5",
        L"SHA1",
        L"SHA256",
        L"SHA384"
    };
    // Static array containing string description of various DH algorithms
    static const wchar_t* const dhString[] =
    {
        L"None",
        L"DH_1",
        L"DH_2",
        L"DH_14",
        L"DH_2048",
        L"DH_ECP_256",
        L"DH_ECP_384",
        L"DH_24"
    };
    // Static array containing string description of various authentication 
    // methods
    static const wchar_t* const authString[] =
    {
        L"Preshared Key",
        L"Certificate",
        L"Kerberos",
        L"Anonymous",
        L"SSL",
        L"NTLMv2",
        L"CGA",
        L"Certificate ECDSA P265",
        L"Certificate ECDSA P384",
        L"SSL ECDSA P265",
        L"SSL ECDSA P384",
        L"EAP",
        L"Reserved"
    };

    //-----------------------------------------
    // Print various pieces of information from the MM SA

    wprintf(L"IKEMM: ");

    // Print MM SA hash, encryption algorithms
    if((mmSa->ikeProposal.integrityAlgorithm.algoIdentifier < 
          RTL_NUMBER_OF(hashString)) &&
       (mmSa->ikeProposal.cipherAlgorithm.algoIdentifier < 
          RTL_NUMBER_OF(cipherString)))
    {
       wprintf(L"%s, %s", 
          cipherString[mmSa->ikeProposal.cipherAlgorithm.algoIdentifier],
          hashString[mmSa->ikeProposal.integrityAlgorithm.algoIdentifier]);
       bLogged = TRUE;
    }

    // Print MM SA DH algorithm
    if(mmSa->ikeProposal.dhGroup < RTL_NUMBER_OF(dhString))
    {
        if(TRUE == bLogged)
        {
            wprintf(L", %s", dhString[mmSa->ikeProposal.dhGroup]);
        }
        else
        {
            wprintf(L"%s", dhString[mmSa->ikeProposal.dhGroup]);
            bLogged = TRUE;
        }
    }

    // Print MM SA auth method
    if(mmSa->ikeCredentials.credentials[0].localCredentials.authenticationMethodType < 
       RTL_NUMBER_OF(authString))
    {
        if(TRUE == bLogged)
        {
            wprintf(L", %s", 
              authString[
                 mmSa->ikeCredentials.credentials[0].localCredentials.authenticationMethodType]
            );
        }
        else
        {
            wprintf(L"%s", 
              authString[
                 mmSa->ikeCredentials.credentials[0].localCredentials.authenticationMethodType]);
            bLogged = TRUE;
        }

        if (IKEEXT_PRESHARED_KEY == 
                mmSa->ikeCredentials.credentials[0].localCredentials.authenticationMethodType)
            DisplayPresharedKeyInfo(
                mmSa->ikeCredentials.credentials[0].localCredentials.presharedKey);

        if (IKEEXT_CERTIFICATE ==
                mmSa->ikeCredentials.credentials[0].localCredentials.authenticationMethodType ||
            IKEEXT_SSL == 
                mmSa->ikeCredentials.credentials[0].localCredentials.authenticationMethodType)
        {
            DisplayCertInfo(
                mmSa->ikeCredentials.credentials[0].localCredentials.certificate);
        }
    }

    // Print EM auth method (if it exists)
    if((mmSa->ikeCredentials.numCredentials > 1) &&
       (mmSa->ikeCredentials.credentials[1].localCredentials.authenticationMethodType < 
       RTL_NUMBER_OF(authString)))
    {
        if(TRUE == bLogged)
        {
            wprintf(L", %s", 
                authString[
                mmSa->ikeCredentials.credentials[1].localCredentials.authenticationMethodType]);
        }
        else
        {
            wprintf(L"%s", 
                authString[
                mmSa->ikeCredentials.credentials[1].localCredentials.authenticationMethodType]);
            bLogged = TRUE;
        }
    }

    if(TRUE == bLogged)
    {
        wprintf(L", %lds", mmSa->ikeProposal.maxLifetimeSeconds);
    }
    else
    {
        wprintf(L"%lds", mmSa->ikeProposal.maxLifetimeSeconds);
    }

    //at this moment bLogged is true.
    switch(mmSa->ikeTraffic.ipVersion)
    {
    case FWP_IP_VERSION_V4:
        DisplayIPv4AddressPair(
            mmSa->ikeTraffic.localV4Address,
            mmSa->ikeTraffic.remoteV4Address);
         break;
    case FWP_IP_VERSION_V6:
        DisplayIPv6AddressPair(
            mmSa->ikeTraffic.localV6Address,
            mmSa->ikeTraffic.remoteV6Address);
        break;
    default:
        break;
    }

    wprintf(L", icookie=0x%" PRIx64, mmSa->cookiePair.initiator);
    wprintf(L", rcookie=0x%" PRIx64, mmSa->cookiePair.responder);
    wprintf(L"\n");
}

//
// Determine if the specified quick-mode security association matches the 
// specified filter and server.
// 
BOOL
IsQmSAExactMatch(
   __in const IPSEC_SA_CONTEXT0* qmSa,
   __in const FWPM_FILTER0* matchingFwpFilter,
   __in const wchar_t* serverSPN,
   __in UINT16 peerPort)
{
    UNREFERENCED_PARAMETER(serverSPN);

    if(!qmSa->outboundSa)
    {
       return FALSE;
    }
    if(qmSa->outboundSa->transportFilter->filterId != 
       matchingFwpFilter->filterId)
    {
       // The filter IDs dont match.
       return FALSE;
    }
    if(qmSa->outboundSa->traffic.remotePort && 
       (qmSa->outboundSa->traffic.remotePort != peerPort))
    {
       // The remote ports dont match.
       return FALSE;
    }

    // SA matches our criteria
    return TRUE;
}

//
// Finds the IPsec SAs (MM & QM) that match a given connection 
// 5-tuple (Local address, Remote address, Local port, Remote port, TCP)
//
DWORD
MatchIPsecSAsForConnection(
   __in const SOCKADDR_STORAGE* localAddress,
   __in const SOCKADDR_STORAGE* peerAddress,
   __in const wchar_t* serverSPN)
{
    DWORD status = 0;
    FWPM_FILTER_CONDITION0 filterConditions[5] = {0};
    HANDLE fwpHandle = NULL;
    FWPM_FILTER_ENUM_TEMPLATE0 filterEnumTemplate = {0};
    HANDLE filterEnumHandle = NULL;
    FWPM_FILTER0** candidateFilters=NULL;
    UINT32 iFilter = 0;
    UINT32 numFiltersReturned=0;
    UINT32 numSasReturned = 0;
    IPSEC_SA_CONTEXT_ENUM_TEMPLATE0 saEnumTemplate;
    HANDLE saEnumHandle = NULL;
    IPSEC_SA_CONTEXT0** qmSaList = NULL;
    IPSEC_SA_CONTEXT0* matchingQmSa = NULL;
    IKEEXT_SA_DETAILS0* matchingMmSa = NULL;
    UINT32 i = 0;
    UINT16 localPort = 0, remotePort = 0;

    ZeroMemory(&saEnumTemplate, sizeof(saEnumTemplate));

    __try
    {
        //-----------------------------------------
        // Open a session handle to FWP
        TRY_DWORD(FwpmEngineOpen0(
            NULL,
            RPC_C_AUTHN_WINNT,
            NULL,
            NULL,
            &fwpHandle));

        //-----------------------------------------
        // Construct the FWP outbound IPsec transport layer filter enum template
        if(localAddress->ss_family == AF_INET)
        {
            filterConditions[3].fieldKey = FWPM_CONDITION_IP_LOCAL_ADDRESS;
            filterConditions[3].matchType = FWP_MATCH_EQUAL;
            filterConditions[3].conditionValue.type = FWP_UINT32;
            filterConditions[3].conditionValue.uint32 = 
               ntohl(((struct sockaddr_in*)localAddress)->sin_addr.s_addr);

            filterConditions[4].fieldKey = FWPM_CONDITION_IP_REMOTE_ADDRESS;
            filterConditions[4].matchType = FWP_MATCH_EQUAL;
            filterConditions[4].conditionValue.type = FWP_UINT32;
            filterConditions[4].conditionValue.uint32 = 
               ntohl(((struct sockaddr_in*)peerAddress)->sin_addr.s_addr);

            localPort = ntohs(((struct sockaddr_in*)localAddress)->sin_port);
            remotePort = ntohs(((struct sockaddr_in*)peerAddress)->sin_port);
            filterEnumTemplate.layerKey = FWPM_LAYER_OUTBOUND_TRANSPORT_V4;
        }
        else
        {
            filterConditions[3].fieldKey = FWPM_CONDITION_IP_LOCAL_ADDRESS;
            filterConditions[3].matchType = FWP_MATCH_EQUAL;
            filterConditions[3].conditionValue.type = FWP_BYTE_ARRAY16_TYPE;
            filterConditions[3].conditionValue.byteArray16 = 
               (FWP_BYTE_ARRAY16*)&(((struct sockaddr_in6*)localAddress)->sin6_addr);

            filterConditions[4].fieldKey = FWPM_CONDITION_IP_REMOTE_ADDRESS;
            filterConditions[4].matchType = FWP_MATCH_EQUAL;
            filterConditions[4].conditionValue.type = FWP_BYTE_ARRAY16_TYPE;
            filterConditions[4].conditionValue.byteArray16 = 
               (FWP_BYTE_ARRAY16*)&(((struct sockaddr_in6*)peerAddress)->sin6_addr);

            localPort = ntohs(((struct sockaddr_in6*)localAddress)->sin6_port);
            remotePort = ntohs(((struct sockaddr_in6*)peerAddress)->sin6_port);
            filterEnumTemplate.layerKey = FWPM_LAYER_OUTBOUND_TRANSPORT_V6;
        }

        filterConditions[0].fieldKey = FWPM_CONDITION_IP_LOCAL_PORT;
        filterConditions[0].matchType = FWP_MATCH_EQUAL;
        filterConditions[0].conditionValue.type = FWP_UINT16;
        filterConditions[0].conditionValue.uint16 = localPort;

        filterConditions[1].fieldKey = FWPM_CONDITION_IP_REMOTE_PORT;
        filterConditions[1].matchType = FWP_MATCH_EQUAL;
        filterConditions[1].conditionValue.type = FWP_UINT16;
        filterConditions[1].conditionValue.uint16 = remotePort;

        filterConditions[2].fieldKey = FWPM_CONDITION_IP_PROTOCOL;
        filterConditions[2].matchType = FWP_MATCH_EQUAL;
        filterConditions[2].conditionValue.type = FWP_UINT8;
        filterConditions[2].conditionValue.uint8 = IPPROTO_TCP;

        /*
        filterConditions[3].fieldKey = FWPM_CONDITION_IP_LOCAL_ADDRESS_TYPE;
        filterConditions[3].matchType = FWP_MATCH_EQUAL;
        filterConditions[3].conditionValue.type = FWP_UINT8;
        filterConditions[3].conditionValue.uint8 = NlatUnicast;*/

        filterEnumTemplate.numFilterConditions = 
            RTL_NUMBER_OF(filterConditions);
        filterEnumTemplate.enumType = FWP_FILTER_ENUM_OVERLAPPING /*FWP_FILTER_ENUM_FULLY_CONTAINED*/;
        filterEnumTemplate.flags= FWP_FILTER_ENUM_FLAG_SORTED /*FWP_FILTER_ENUM_FLAG_BEST_TERMINATING_MATCH*/;
        filterEnumTemplate.filterCondition = filterConditions;
        filterEnumTemplate.actionMask = FWP_ACTION_FLAG_TERMINATING;

        //-----------------------------------------
        // Enumerate FWP outbound IPsec transport layer filters
        TRY_DWORD(FwpmFilterCreateEnumHandle0(
            fwpHandle,
            &filterEnumTemplate,
            &filterEnumHandle));

        TRY_DWORD(FwpmFilterEnum0(
            fwpHandle,
            filterEnumHandle,
            cCANDIDATE_FILTERS_REQUESTED,
            &candidateFilters,
            &numFiltersReturned));

        //-----------------------------------------
        // Enumerate the IPsec QM SAs
        TRY_DWORD(IPsecSaContextCreateEnumHandle0(
            fwpHandle,
            &saEnumTemplate,
            &saEnumHandle));

        TRY_DWORD(IPsecSaContextEnum0(
            fwpHandle,
            saEnumHandle,
            UINT_MAX,
            &qmSaList,
            &numSasReturned));
        if (numSasReturned < 1)
        {
            status = ERROR_INVALID_PARAMETER;
            wprintf(L"IPsecSaContextEnum0 didn't return any SAs\n");
            __leave;
        }

        //
        // Now look for a matching pair: transport layer filter + QM SA
        //

        for( ; 
            iFilter < numFiltersReturned && NULL == matchingQmSa; 
            iFilter++)
        {
            // Check for a matching filter
            if((candidateFilters[iFilter]->action.type & FWP_ACTION_FLAG_CALLOUT) &&
               ((memcmp(
                     &candidateFilters[iFilter]->action.calloutKey,
                     &FWPM_CALLOUT_IPSEC_OUTBOUND_TRANSPORT_V4,
                     sizeof(GUID)
                  ) == 0) ||
                 (memcmp(
                     &candidateFilters[iFilter]->action.calloutKey,
                     &FWPM_CALLOUT_IPSEC_OUTBOUND_TUNNEL_V4,
                     sizeof(GUID)
                  ) == 0) ||
                 (memcmp(
                     &candidateFilters[iFilter]->action.calloutKey,
                     &FWPM_CALLOUT_IPSEC_OUTBOUND_TRANSPORT_V6,
                     sizeof(GUID)
                  ) == 0) ||
                 (memcmp(
                     &candidateFilters[iFilter]->action.calloutKey,
                     &FWPM_CALLOUT_IPSEC_OUTBOUND_TUNNEL_V6,
                     sizeof(GUID)
                  ) == 0)))
            {
                //-----------------------------------------
                // Go through the enumerated QM SAs and look for an exact match
                for (i = 0; i < numSasReturned; i++)
                {
                    if (IsQmSAExactMatch(
                        qmSaList[i],
                        candidateFilters[iFilter],
                        serverSPN,
                        remotePort))
                    {
                        matchingQmSa = qmSaList[i];
                        break;
                    }
                }
            }
        }

        if(iFilter == numFiltersReturned)
        {
            status = ERROR_INVALID_PARAMETER;
            wprintf(L"FwpmFilterEnum0 didn't return any matching filters\n");
            __leave;
        }

        if(!matchingQmSa)
        {
           status = ERROR_INVALID_PARAMETER;
           wprintf(L"Didn't find any matching SAs\n");
           __leave;
        }

        //-----------------------------------------
        // Get the IPsec MM SA using the MM SA ID
        TRY_DWORD(IkeextSaGetById0(
            fwpHandle,
            matchingQmSa->outboundSa->saBundle.mmSaId,
            &matchingMmSa));

        //-----------------------------------------
        // Print the SAs
        PrintMmSa(matchingMmSa);
        PrintQmSa(matchingQmSa);
    }
    __finally
    {
        if(candidateFilters)
           FwpmFreeMemory0((void **) &candidateFilters);
        if(qmSaList)
           FwpmFreeMemory0((void **) &qmSaList);
        if(filterEnumHandle)
           FwpmFilterDestroyEnumHandle0(fwpHandle, filterEnumHandle);
        if(saEnumHandle)
           IPsecSaContextDestroyEnumHandle0(fwpHandle, saEnumHandle);
        if(fwpHandle)
           FwpmEngineClose0(fwpHandle);
    }
    return status;
}

PSOCKET_ADDRESS_LIST 
CreateRemoteAddressList(
    __in const wchar_t* serverHostDnsName,
    __in const wchar_t* serverPort,
    __in const int addrFamily)
{
    DWORD dwSize = 0;
    DWORD i = 0;
    DWORD dwAddrCount = 0;
    DWORD status = 0;
    const addrinfoW *ptr = NULL;
    void *pAddr = NULL;
    DWORD dwAddrSize = 0;
    struct addrinfoW aiHints = {0};
    struct addrinfoW* aiList = NULL;
    PSOCKET_ADDRESS_LIST socketAddressList = NULL;
    BOOL bSuccess = FALSE;

    __try
    {
        aiHints.ai_family = addrFamily;
        aiHints.ai_socktype = SOCK_STREAM;
        aiHints.ai_protocol = IPPROTO_TCP;

        status = GetAddrInfo(
            serverHostDnsName,
            serverPort,
            &aiHints,
            &aiList);
        if (status)
        {
           wprintf(L"GetAddrInfo returned error %ld\n", status);
           __leave;
        }

        ptr = aiList;
        while(NULL != ptr)
        {
            ++dwAddrCount;
            ptr = ptr->ai_next;
        }

        dwSize = SIZEOF_SOCKET_ADDRESS_LIST(dwAddrCount) + 
                 (dwAddrCount * sizeof(SOCKADDR_STORAGE));

        TRY_ALLOC(socketAddressList = 
            (PSOCKET_ADDRESS_LIST) Alloc(dwSize));

        dwAddrSize = AF_INET == aiList->ai_family ? 
            sizeof(SOCKADDR_IN) : sizeof(SOCKADDR_IN6);

        for (ptr = aiList, i = 0; NULL != ptr; ptr = ptr->ai_next, i++)
        {
            TRY_ALLOC(pAddr = Alloc(dwAddrSize));
            CopyMemory(pAddr,ptr->ai_addr, dwAddrSize);
            SS_PORT((SOCKADDR*)pAddr) = htons((USHORT)_wtoi(serverPort));
            socketAddressList->Address[i].lpSockaddr = (SOCKADDR*)pAddr;
            socketAddressList->Address[i].iSockaddrLength = (int)dwAddrSize;
        }
        socketAddressList->iAddressCount = dwAddrCount;
        bSuccess = TRUE;
    }
    __finally
    {
        if(aiList)
            FreeAddrInfo(aiList);

        if(FALSE == bSuccess)
        {
            if(socketAddressList)
            {
                for(i = 0; i < dwAddrCount; i++)
                {
                    if(socketAddressList->Address[i].lpSockaddr)
                        Free(socketAddressList->Address[i].lpSockaddr);
                }
                Free(socketAddressList);
            }
        }
    }
    return socketAddressList;
}

SOCKET 
CreateSocket(
    __in_opt    const wchar_t* serverSPN,
    __in        const int addrFamily,
    __in        const BOOL fVerbose,
    __in        const BOOL fSPN)
{
    DWORD status = ERROR_SUCCESS;
    SOCKET sock = INVALID_SOCKET;
    SOCKET_SECURITY_SETTINGS securitySettings;
    SOCKET_PEER_TARGET_NAME* peerTargetName = NULL;
    DWORD serverSpnStringLen = (DWORD) wcslen(serverSPN);
    DWORD peerTargetNameLen = sizeof(SOCKET_PEER_TARGET_NAME) + 
        (serverSpnStringLen * sizeof(wchar_t));
    wchar_t* wcsAddrFamily = NULL;
    INT off = 0;

    ZeroMemory(&securitySettings, sizeof(securitySettings));

    __try
    {
        // Create a TCP socket
        sock = WSASocket(
            addrFamily,
            SOCK_STREAM,
            IPPROTO_TCP,
            NULL,
            0,
            0);

        if(TRUE == fVerbose)
        {
            switch(addrFamily)
            {
            case AF_INET:
                wcsAddrFamily = L"AF_INET";
                break;
            case AF_INET6:
                wcsAddrFamily = L"AF_INET6";
                break;
            case AF_UNSPEC:
                wcsAddrFamily = L"AF_UNSPEC";
                break;
            default:
                break;
            }
            wprintf(L"Create socket <%s, SOCK_STREAM, IPPROTO_TCP>, return code=%ld\n",
                wcsAddrFamily, 
                WSAGetLastError());
        }

        if (INVALID_SOCKET == sock)
        {
            status = WSAGetLastError();
            wprintf(L"WSASocket returned error %ld\n", status); 
            __leave;
        }

        if(AF_INET6 == addrFamily)
        {
            TRY_WSACALL(setsockopt(
                sock,
                IPPROTO_IPV6,
                IPV6_V6ONLY,
                (char*)&off,
                sizeof(off)));
        }

        //-----------------------------------------
        // Turn on security for the socket.
        securitySettings.SecurityProtocol = SOCKET_SECURITY_PROTOCOL_DEFAULT;
        securitySettings.SecurityFlags = SOCKET_SETTINGS_GUARANTEE_ENCRYPTION;
        TRY_WSACALL(WSASetSocketSecurity (
            sock,
            &securitySettings,
            sizeof(securitySettings),
            NULL,
            NULL));

        if (TRUE == fSPN)
        {
            //-----------------------------------------
            // Specify the server SPN
            TRY_ALLOC(peerTargetName = (SOCKET_PEER_TARGET_NAME *) Alloc(
                sizeof(wchar_t) + peerTargetNameLen));
            peerTargetName->SecurityProtocol = securitySettings.SecurityProtocol;
            peerTargetName->PeerTargetNameStringLen = serverSpnStringLen;

            RtlCopyMemory(
                (BYTE*)peerTargetName->AllStrings,
                (BYTE*)serverSPN,
                serverSpnStringLen * sizeof(wchar_t));

            TRY_WSACALL(WSASetSocketPeerTargetName(
                sock,
                peerTargetName,
                peerTargetNameLen,
                NULL,
                NULL));
        }
    }
    __finally
    {
        if(peerTargetName)
            Free(peerTargetName);
        if(ERROR_SUCCESS != status)
            if(INVALID_SOCKET != sock)
                closesocket(sock);
    }
    return sock;
}

DWORD
WINAPI
SecureTcpConnect(
    __in const wchar_t* serverHostDnsName,
    __in const wchar_t* serverPort,
    __in const int addrFamily,
    __in const BOOL fVerbose,
    __in const BOOL fSPN)
{
    DWORD status = ERROR_SUCCESS;
    wchar_t serverSPN [MAX_PATH] = {0};
    DWORD serverSPNLen = MAX_PATH;
    PSOCKET_ADDRESS_LIST addressList = NULL;
    SOCKET sock = INVALID_SOCKET;
    SOCKADDR_STORAGE localAddr = {0};
    int localAddrLen = sizeof(localAddr);
    int sockErr = 0;
    wchar_t localIPAddr [MAX_PATH];
    wchar_t remoteIPAddr [MAX_PATH];
    DWORD cb = 0;
    DWORD dwTickCount1 = 0;
    DWORD dwTickCount2 = 0;
    const float fMsInASecond = 1000.0;

    __try
    {
        dwTickCount1 = GetTickCount();

        if (TRUE == fSPN)
        {
            // Get the server SPN from the DNS name.
            TRY_DWORD(DsMakeSpn(
                L"host",
                serverHostDnsName,
                NULL,
                0,
                NULL,
                &serverSPNLen,
                serverSPN));
        }

        // Create the socket
        sock = CreateSocket(serverSPN, addrFamily, fVerbose, fSPN);
        if(INVALID_SOCKET == sock)
        {
            status = WSAGetLastError();
            status = ERROR_SUCCESS == status ? SOCKET_ERROR : status;
            __leave;
        }

        //----------------------------------------
        // Prepare the destination address list
        addressList = CreateRemoteAddressList(
            serverHostDnsName,
            serverPort,
            addrFamily);
        if(NULL == addressList)
        {
            status = WSAGetLastError();
            status = ERROR_SUCCESS == status ? SOCKET_ERROR : status;
            wprintf(L"No valid address available.\n");
           __leave;
        }

        // Output target info
        cb = sizeof(remoteIPAddr) / sizeof(remoteIPAddr [0]);
        TRY_DWORD(WSAAddressToString(
            addressList->Address[0].lpSockaddr,
            addressList->Address[0].iSockaddrLength,
            NULL,
            remoteIPAddr,
            &cb));

        wprintf(L"TCP over IPsec connect to dest: %s, addr %s...\n",
            serverHostDnsName, remoteIPAddr);

        //-----------------------------------------
        // Connect to the server
        sockErr = WSAConnect(
            sock,
            addressList->Address[0].lpSockaddr,
            addressList->Address[0].iSockaddrLength,
            NULL,
            NULL,
            NULL,
            NULL);
        if (sockErr == SOCKET_ERROR)
        {
            status = WSAGetLastError();
            wprintf(L"WSAConnect returned error %ld\n", status);
            __leave;
        }

        //-----------------------------------------
        // Get local address bound to the socket   
        sockErr = getsockname(
            sock,
            (struct sockaddr*)&localAddr,
            &localAddrLen);
        if (sockErr == SOCKET_ERROR)
        {
           status = WSAGetLastError();
           wprintf(L"getsockname returned error %ld\n", status);
           __leave;
        }

        // Output connection information
        cb = sizeof(localIPAddr) / sizeof(localIPAddr [0]);
        TRY_DWORD(WSAAddressToString(
            (LPSOCKADDR)&localAddr,
            localAddrLen,
            NULL,
            localIPAddr,
            &cb));

        dwTickCount2 = GetTickCount();
        wprintf(L"TCP over IPSec WSAConnect: success, %.2fs from %s to %s\n",
            (dwTickCount2 - dwTickCount1)/fMsInASecond, localIPAddr, remoteIPAddr);

        if(TRUE == fVerbose)
        {
            //-----------------------------------------
            // Match and print IPSec SA information for the connection 
            TRY_DWORD(MatchIPsecSAsForConnection(
                &localAddr, 
                (SOCKADDR_STORAGE*)addressList->Address[0].lpSockaddr,
                serverSPN));
        }
    }
    __finally
    {
        if(addressList)
        {
            for(int i = 0; i < addressList->iAddressCount; i++)
            {
                if(addressList->Address[i].lpSockaddr)
                    Free(addressList->Address[i].lpSockaddr);
            }
            Free(addressList);
        }

        if(TRUE == fVerbose && ERROR_SUCCESS == status)
        {
            wprintf(L"Disconnecting, IKE MM and QM deleted\n");
        }
        if(INVALID_SOCKET != sock)
            closesocket(sock);
    }
    
    return status;
}

//
// Display usage information to console
// 
void
WINAPI
Usage(void)
{
    wprintf(L"IPsecPing Version 1.7 Copyright 2016 JW Secure, Inc.\n");
    wprintf(L"\nUsage:\n");
    wprintf(L" ipsecping.exe [-s] [-v] [-4 | -6] [-p <port>] <DNSname>\n");
    wprintf(L" -v: verbose, default is off\n");
    wprintf(L" -4: for IPv4 only, the default value\n");
    wprintf(L" -6: for IPv6 only\n");
    wprintf(L" -p: connect-to port, default is 135\n");
    wprintf(L" -s: set a Service Principal Name (requires Kerberos)\n");
}

//
// Main application entry point
// 
int wmain(
    __in              int argc, 
    __in_ecount(argc) WCHAR* argv[])
{
    DWORD status = ERROR_SUCCESS;
    BOOL fVerbose = FALSE;
    BOOL fSPN = FALSE;
    BOOL fInvalidParam = FALSE;
    LPWSTR wszError = NULL;
    WORD version = MAKEWORD(2, 2);
    BOOL wsaCleanup = FALSE;
    wchar_t* serverHostDnsName = NULL;
    wchar_t* serverPort = L"135";
    WSADATA wsaData = {0};
    int addrFamily = AF_INET;

    __try
    {
        --argc;
        ++argv;

        if (argc)
        {
            while (argc)
            {
                if (L'/' != **argv && L'-' != **argv)
                {
                    // 
                    // Doesn't look like a command-line switch
                    // For now, assume it's a destination DNS name
                    // 
                    serverHostDnsName = *argv;
                }
                else
                {
                    switch (argv [0][1])
                    {
                    case L'v':
                        fVerbose = TRUE;
                        break;

                    case L'4':
                        addrFamily = AF_INET;
                        break;

                    case L'6':
                        addrFamily = AF_INET6;
                        break;

                    case L'p':
                        --argc;
                        ++argv;
                        serverPort = *argv;
                        break;

                    case L's':
                        fSPN = TRUE;
                        break;

                    default:
                        fInvalidParam = TRUE;
                        __leave;
                    }
                }

                --argc;
                ++argv;
            }
        }
        else
        {
            fInvalidParam = TRUE;
            __leave;
        }

        status = WSAStartup(version, &wsaData);
        if (status)
        {
           wprintf(L"WSAStartup returned error %ld\n", status);
           __leave;
        }

        wsaCleanup = TRUE;

        // Establish secure TCP connection to the server
        TRY_DWORD(SecureTcpConnect(
            serverHostDnsName, 
            serverPort, 
            addrFamily, 
            fVerbose,
            fSPN));
    }
    __finally
    {
        if(wsaCleanup)
           WSACleanup();
    }

    if (TRUE == fInvalidParam)
        Usage();
    else if (ERROR_SUCCESS != status)
    {
        if (0 != FormatMessage(
            FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
            NULL,
            status,
            MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
            (LPWSTR) &wszError,
            0,
            NULL))
        {
            wprintf(L" failed\n0x%x - %s", status, wszError);
            if (NULL != wszError)
                LocalFree(wszError);
        }
        else
        {
            wprintf(L" failed\n0x%x\n", status);
        }
    }
   
    return status;
}